from collections import Counter
import random

def balance_dataset(input_file, output_file, max_per_class, exclude_classes):
    with open(input_file, 'r') as f:
        lines = f.readlines()
    
    # 解析文件中的图像名称和标签
    data = [line.strip().split() for line in lines]
    labels = [int(item[1]) for item in data]
    
    # 统计每个类别的数量
    label_counts = Counter(labels)
    print("Original label counts:", label_counts)
    
    # 按类别分组数据
    label_groups = {label: [] for label in label_counts}
    for item in data:
        label_groups[int(item[1])].append(item)
    
    # 对类别进行平衡并排除指定类别
    balanced_data = []
    for label, items in label_groups.items():
        if label in exclude_classes:
            print(f"Excluding class {label}")
            continue  # 跳过需要排除的类别
        if len(items) > max_per_class:
            balanced_data.extend(random.sample(items, max_per_class))
        else:
            balanced_data.extend(items)
    
    # 写入新的文件
    with open(output_file, 'w') as f:
        for item in balanced_data:
            f.write(' '.join(item) + '\n')
    
    # 打印新数据集的类别统计信息
    new_label_counts = Counter([int(item[1]) for item in balanced_data])
    print("Balanced label counts:", new_label_counts)

# 输入和输出文件路径
input_file = 'plots/tsne/train.txt'  # 原始文件
output_file = 'plots/tsne/test.txt'  # 平衡后的文件

# 每个类别的最大样本数
max_per_class = 1500

# 要排除的类别
# exclude_classes = [1, 4]
exclude_classes = []

# 调用函数
balance_dataset(input_file, output_file, max_per_class, exclude_classes)
